package com.otr.ai.web;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentReader;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import java.util.List;

@Slf4j
@RestController
@RequestMapping("/upload")
@RequiredArgsConstructor
public class VectorStoreController {

  private final VectorStore vectorStore;

  /**
   * 处理PDF文档的解析、分割和嵌入存储。 使用 PagePdfDocumentReader 解析PDF文档并生成 Document 列表。 使用
   * TokenTextSplitter将文档分割成更小的部分。 将分割后的文档添加到向量存储中，以便后续检索和生成。
   */
  @PostMapping("/pdf")
  public String upload(@RequestParam("file") MultipartFile file) {

    // 1. parse document
    DocumentReader reader = new PagePdfDocumentReader(file.getResource());
    List<Document> documents = reader.get();
    log.info("{} documents loaded", documents.size());

    // 2. split trunks
    TokenTextSplitter tokenTextSplitter = new TokenTextSplitter();
    List<Document> splitDocuments = tokenTextSplitter.apply(documents);
    log.info("{} documents split", splitDocuments.size());

    // 3. create embedding and store to vector store
    log.info("create embedding and save to vector store");
    vectorStore.add(splitDocuments);
    return "ok";
  }
}
